# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/3/21 13:40
# @Function: 第一步：将2010land30m合并成31个省；将每个省替换代码，替换代码为新国标
"""
E:\land30m2010\anhui

"""

import os
import numpy as np
import rasterio
from rasterio.merge import merge

def select_valid_pixels(merged_data, new_data, merged_mask, new_mask, **kwargs):
    """
    Returns the valid pixel values from both data sources.
    If both sources have valid pixels in the same position,
    the pixel from new_data is chosen.
    """
    # 创建一个布尔掩码，表示merged_mask和new_mask都为True的位置
    valid_mask = np.logical_and(merged_mask, new_mask)

    # 创建一个临时掩码，用于标记merged_data中需要被更新的像素位置
    update_mask = np.copy(valid_mask)

    # 如果merged_mask为True但new_mask为False，保留merged_data的值
    np.logical_and(merged_mask, np.logical_not(new_mask), out=valid_mask)
    np.logical_or(update_mask, valid_mask, out=update_mask)

    # 如果new_mask为True但merged_mask为False，更新merged_data的值为new_data
    np.logical_and(np.logical_not(merged_mask), new_mask, out=valid_mask)
    np.logical_or(update_mask, valid_mask, out=update_mask)

    # 更新merged_data的值
    np.copyto(merged_data, new_data, where=update_mask, casting="unsafe")

def merge_tifs_in_directory(province_dir: str, output_file: str) -> None:
    """
    合并指定目录中的所有 TIFF 文件，并将结果写入到输出文件中。

    参数:
    province_dir (str): 包含 TIFF 文件的目录路径。
    output_file (str): 合并结果的输出文件路径。
    """
    # 使用 rasterio.open 读取 TIFF 文件
    rasters = [rasterio.open(os.path.join(province_dir, f)) for f in os.listdir(province_dir) if f.endswith('.tif')]

    if not rasters:
        print("没有找到有效的 TIFF 文件")
        return

    # 合并 TIFF 文件
    merged_raster, out_trans = merge(rasters, nodata=0, method=select_valid_pixels)

    # 写入合并结果到输出文件
    with rasterio.open(
            output_file,
            'w',
            driver='GTiff',
            height=merged_raster.shape[1],
            width=merged_raster.shape[2],
            count=merged_raster.shape[0],
            dtype=np.uint16,  # 设置数据类型为 uint8
            crs='EPSG:4326',
            transform=out_trans,
            nodata=0,
            compress='lzw'
    ) as dst:
        dst.write(merged_raster)



province_dir = r'H:\land30m2020\replace_output\GLC_FCS30_2020_replace_beast'
output_file = r'H:\land30m2020\replace_output\GLC_FCS30_2020_replace_beast\merged\GLC_FCS30_2020_replace_beast.tif'

merge_tifs_in_directory(province_dir, output_file)
print(f"合并完成:{output_file}")

province_dir = r'H:\land30m2020\replace_output\GLC_FCS30_2020_replace'
output_file = r'H:\land30m2020\replace_output\GLC_FCS30_2020_replace\merged\GLC_FCS30_2020_replace.tif'

merge_tifs_in_directory(province_dir, output_file)
print(f"合并完成:{output_file}")
